import java.util.ArrayList;
import java.util.List;

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int numRows = nums.size();
        int maxCols = 0;

        // 找到最大列数
        for (List<Integer> row : nums) {
            maxCols = Math.max(maxCols, row.size());
        }

        // 创建结果数组
        int[] diagonalOrder = new int[numRows * maxCols];
        int index = 0;

        // 遍历对角线
        for (int sum = 0; sum < numRows + maxCols - 1; sum++) {
            List<Integer> currentDiagonal = new ArrayList<>();
            for (int i = 0; i <= sum; i++) {
                int j = sum - i;
                if (i < numRows && j < nums.get(i).size()) {
                    currentDiagonal.add(nums.get(i).get(j));
                }
            }
            if (sum % 2 == 0) {
                // 偶数对角线，按顺序添加当前对角线上的元素
                for (int num : currentDiagonal) {
                    diagonalOrder[index++] = num;
                }
            } else {
                // 奇数对角线，反转当前对角线上的元素
                for (int i = currentDiagonal.size() - 1; i >= 0; i--) {
                    diagonalOrder[index++] = currentDiagonal.get(i);
                }
            }
        }

        return diagonalOrder;
    }
}
